import os
import sys
import click
from alembic import command
from alembic.config import Config
from loguru import logger

# 获取项目根目录
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(BASE_DIR)

# 获取alembic.ini路径
alembic_cfg = Config(os.path.join(BASE_DIR, "alembic.ini"))

@click.group()
def cli():
    """数据库迁移管理工具"""
    pass

@cli.command()
def init():
    """初始化数据库"""
    try:
        from app.core.init_db import init_db
        logger.info("开始初始化数据库...")
        init_db()
        logger.info("数据库初始化完成")
    except Exception as e:
        logger.error(f"数据库初始化失败: {str(e)}")
        sys.exit(1)

@cli.command()
def migrate():
    """生成新的迁移脚本"""
    try:
        logger.info("生成迁移脚本...")
        command.revision(alembic_cfg, autogenerate=True)
        logger.info("迁移脚本生成完成")
    except Exception as e:
        logger.error(f"生成迁移脚本失败: {str(e)}")
        sys.exit(1)

@cli.command()
def upgrade():
    """升级数据库到最新版本"""
    try:
        logger.info("开始升级数据库...")
        command.upgrade(alembic_cfg, "head")
        logger.info("数据库升级完成")
    except Exception as e:
        logger.error(f"数据库升级失败: {str(e)}")
        sys.exit(1)

@cli.command()
@click.argument("revision", default="-1")
def downgrade(revision):
    """回滚数据库到指定版本"""
    try:
        logger.info(f"开始回滚数据库到版本 {revision}...")
        command.downgrade(alembic_cfg, revision)
        logger.info("数据库回滚完成")
    except Exception as e:
        logger.error(f"数据库回滚失败: {str(e)}")
        sys.exit(1)

@cli.command()
def history():
    """查看迁移历史"""
    try:
        logger.info("迁移历史:")
        command.history(alembic_cfg)
    except Exception as e:
        logger.error(f"获取迁移历史失败: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    cli() 